""" Data model (Torrent, TorrentPeerEntry and TorrentPiece).

Originally part of the open-tracker project (http://code.google.com/p/open-tracker/),
modified when working on torrent-tracker.
"""

__author__ = 'allen@thebends.org (Allen Porter) and gavinandresen@gmail.com (Gavin Andresen)'

import hashlib
import logging

from google.appengine.ext import db
from BTL.bencode import bencode, bdecode

class Torrent(db.Model):
    # info_hash stored as Torrent.key().name()
    bencoded_data = db.BlobProperty()
    name = db.StringProperty()
    date = db.DateTimeProperty(auto_now_add=True)
    complete = db.BooleanProperty(default=False)

    # Total number of times the tracker has registered a completion
    downloaded = db.IntegerProperty()

    # Override __getattr__ to bdecode bencoded_data, cache it, and
    # allow direct access to the keys in it.
    # So you can do stuff like:
    #  Torrent foo = Torrent.get_by_hash(hash)
    #  piece_length = foo.info['piece length']
    def __getattr__(self, attribute):
        if '_bdecoded' not in self.__dict__:
            self.__dict__['_bdecoded'] = bdecode(self.bencoded_data)

        if attribute in self._bdecoded:
            self.__dict__[attribute] = self._bdecoded[attribute]
            return self.__dict__[attribute]

        raise AttributeError

    @staticmethod
    def get_by_hash(hash):
        return Torrent.gql("WHERE __key__ = KEY('Torrent', :1)", hash).get()

    @staticmethod
    def create_from_dict(d):
        torrent_data = bencode(d)
        info_string = bencode(d['info'])
        info_hash = hashlib.sha1(info_string).hexdigest()
        piece_string = d['info']['pieces']
        n_pieces = len(piece_string)/20
        return Torrent(key_name = info_hash,
                       bencoded_data = torrent_data,
                       name = d['info']['name'])

    def get_piece(self, piece_number):
        p = TorrentPiece.gql("WHERE torrent = :1 AND piece_number=:2", self, piece_number).get()
        if p is None:
            return None
        return p.data

    def put_piece(self, piece_number, piece_data):
        piece_hash = hashlib.sha1(piece_data).hexdigest()
        if piece_hash.decode('hex_codec') not in self.info['pieces']:
            logging.error("ERROR: piece hash mismatch!  Corrupt .torrent, or on-disk files have changed.")
            sys.exit(1)
        existing_piece_key = db.GqlQuery(
            "SELECT __key__ FROM TorrentPiece WHERE torrent = :1 AND piece_number=:2", self, piece_number).get()
        if existing_piece_key is None:
            p = TorrentPiece(torrent=self, data=piece_data, piece_number=piece_number)
            p.put()
            logging.info("Put piece "+piece_hash)
        else:
            logging.info("Skipped piece "+piece_hash)

    def get_data(self, start, length):
        if length < 1:
            return ""
        piece_length = self.info['piece length']
        (piece, piece_offset) = divmod(start, piece_length)
        logging.debug("Looking for range (%d-%d); got piece %d (offset %d, piece_length %d)"%(
                start, length, piece, piece_offset, piece_length))
        data = self.get_piece(piece)
        if data is None:
            return ""
        
        if piece_offset+length <= piece_length:
            return data[piece_offset:piece_offset+length]
        else:
            return (data[piece_offset:] +
                    self.get_data((piece+1)*piece_length, length-(piece_length-piece_offset)))

    # Override to delete piece hashes, too
    def delete(self):
        for piece_key in db.GqlQuery("SELECT __key__ FROM TorrentPiece WHERE torrent = :1", self):
            db.delete(piece_key)
        super(Torrent, self).delete()

class TorrentPiece(db.Model):
    torrent = db.Reference(Torrent, required=True)
    piece_number = db.IntegerProperty(required=True)
    data = db.BlobProperty(required=True)

class TorrentPeerEntry(db.Model):
  torrent = db.Reference(Torrent)

  # IP:port of the client
  ip = db.StringProperty(required=True)
  port = db.IntegerProperty(required=True)

  # Client supplied identifier
  peer_id = db.ByteStringProperty(required=True)

  # Total amount downloaded (probably bytes)
  downloaded = db.IntegerProperty()

  # Total amount uploaded (probably bytes)
  uploaded = db.IntegerProperty()

  # Last time we saw the client
  last_datetime = db.DateTimeProperty(required=True)

